# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import enum
import functools
import hashlib
import json
import typing
import uuid

import requests
import tornado.gen
import tornado.httpclient

import coin.flow.subscriber as flow_subscriber
import coin.flow.topic_map as topic_map
import coin.proto.coin_market_enums_pb2 as coin_enum

from coin.base.config import Config
from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_fill_util import FillManager, LastFill
from coin.exchange.base.order_gateway_ready import SpotOrderGatewayReady
from coin.exchange.base.rest.task_scheduler import Scheduler, wrap_async_task
from coin.exchange.base.rest_order_gateway_base import PrivateExchangeMessage
from coin.exchange.okex.kr_rest.currency import OkexCurrency
from coin.exchange.okex.kr_rest.product import OkexProduct
from coin.exchange.okex_v3.kr_rest.private_client import OkexPrivateParser
from coin.exchange.util.rest_rate_limit import RestRateLimiter
from coin.proto.coin_order_gateway_pb2 import OrderGatewayConfig, OrderEvent
from coin.proto.coin_query_pb2 import CurrencyBalance, ProductFillElement

from coin.exchange.okex_v3.ws.private_subscriber import (
    OkexWsSpotPrivateClient,)
from coin.exchange.okex_v3.kr_rest.native_private_client import (
    OkexNativePrivateClient,
    OkexOrderSide,
    OkexOrderType,
    OkexOrderStatus,
    OkexOrderProperty,
)
from coin.exchange.base.rest_order_gateway_base import (
    RestOrderGateway,
    exchange_message_handler,
    create_order_internal,
    to_proto_order_side,
    BaseBalanceOrderGatewayInfo,
)
from coin.exchange.base.order_gateway import (
    Order,
    OrderSide,
    OrderType,
    OrderGatewayStatus,
)


def create_okex_order_from_order_update(data, timestamp, order_id=None):
  """
  WS:
  {
    "table": "spot/order",
    "data": [
      {
        "filled_notional": "0",
        "filled_size": "0",
        "instrument_id": "BTC-USDT",
        "margin_trading": "1",
        "notional": "",
        "order_id": "2041964978506752",
        "price": "3718.1465",
        "side": "buy",
        "size": "0.01",
        "status": "cancelled",
        "timestamp": "2018-12-27T06:58:45.000Z",
        "type": "limit"
      }
    ]
  }
  20190722
  {
  "table": "spot/order",
  "data": [
    {
      "client_oid": "pl79d66252",
      "filled_notional": "0",
      "filled_size": "0",
      "instrument_id": "BTC-USDT",
      "last_fill_px": "0",
      "last_fill_qty": "0",
      "last_fill_time": "1970-01-01T00:00:00.000Z",
      "margin_trading": "1",
      "notional": "",
      "order_id": "3214894999344128",
      "order_type": "1",
      "price": "10447.9",
      "side": "buy",
      "size": "0.00443175",
      "state": "0",
      "status": "open",
      "timestamp": "2019-07-22T10:30:04.667Z",
      "type": "limit"
    }
  ]
  }
  REST:
  [
    {
      "created_at": "2018-10-24T09:59:15Z",
      "executed_value": "0",
      "filled_notional": "0",
      "filled_size": "0",
      "funds": "",
      "instrument_id": "BTC-USDT",
      "notional": "",
      "order_id": "1680287676308480",
      "price": "5800",
      "product_id": "BTC-USDT",
      "side": "buy",
      "size": "0.002",
      "status": "open",
      "timestamp": "2018-10-24T09:59:15Z",
      "type": "limit"
    },
  ]
  """
  native_symbol = data['instrument_id']
  product = OkexProduct.FromStrNativeProduct(native_symbol)
  exchange_order_id = str(data['order_id'])
  if data['price']:
    price_orig = float(data['price'])
  else:
    price_orig = None
  qty_orig = float(data['size'])
  qty_accum_fill = float(data['filled_size'])
  price_avg_fill = None
  pq_accum_fill = float(data['filled_notional'])
  if pq_accum_fill > 0:
    price_avg_fill = pq_accum_fill / qty_accum_fill

  okex_order_side = OkexOrderSide(data['side'])
  okex_order_type = OkexOrderType(data['type'])
  order_side = getattr(OrderSide, okex_order_side.name)
  order_type = getattr(OrderType, okex_order_type.name)

  status = OkexOrderStatus(data['status'])
  fill = FillManager(price_orig=price_orig,
                     qty_orig=qty_orig,
                     qty_accum_fill=qty_accum_fill,
                     price_avg_fill=price_avg_fill)

  order = Order(order_id=order_id,
                product=product,
                price=price_orig,
                qty=qty_orig,
                order_side=order_side,
                order_type=order_type,
                tag=None,
                accepted=True,
                cancel_sent=False,
                internal=create_order_internal(
                    timestamp=timestamp,
                    exchange_order_id=exchange_order_id,
                    fill=fill,
                    status=status,
                ))
  return order


class OkexOrderGatewayInfo(BaseBalanceOrderGatewayInfo):
  CurrencyType = OkexCurrency
  ProductType = OkexProduct
  BalanceType = float
  PositionType = None

  def __init__(self, currencies: typing.List[OkexCurrency], products: typing.List[OkexProduct]):
    super().__init__(currencies, products)
    for currency in currencies:
      empty_balance = CurrencyBalance(
          currency=currency.currency,
          currency_native=currency.native_currency,
          total=0.,
          available=0.,
          hold=0.,
      )
      self.set_balance(currency, 0., internal=empty_balance)

  def update_balance(self, data):
    account_balance = OkexPrivateParser.parse_balance(data)
    for balance in account_balance.each_balance:
      currency = OkexCurrency.FromStrCurrency(balance.currency)
      if currency not in self.currencies:
        continue
      self.set_balance(currency, balance.total, internal=balance)


class OkexOrderGateway(RestOrderGateway):
  _exchange = 'Okex'
  _market_type = 'Spot'
  _api_version = 'v3'

  OGInfoType = OkexOrderGatewayInfo

  class _MsgType(enum.IntEnum):
    UNKNOWN = 0
    REST_ACCOUNT_BALANCE = 1
    WS_ACCOUNT_BALANCE = 2
    OPEN_ORDERS = 3
    ORDER_UPDATE = 4
    REST_FILLS = 5

  def __init__(
      self,
      currencies: typing.List[OkexCurrency],
      products: typing.List[OkexProduct],
      config: Config,
      og_config=None,
      logger=None,
  ):
    super().__init__(currencies, products, config, og_config, logger)
    self._og_ready = SpotOrderGatewayReady()
    self._og_ready.set_status(OrderGatewayStatus.INITIALIZING)

    self._private_client = None
    self._private_subscriber = None
    self._scheduler = None

    self._ioloop = None
    self._flow_sub = None
    self._publisher = None

    default_og_config = OrderGatewayConfig(
        account_query_period=1,
        orders_snapshot_query_period=5,
        fill_query_period=5,
        use_fill_checker=True,
        rest_lib='pycurl',
    )
    default_og_config.MergeFrom(self._og_config)
    self._og_config = default_og_config
    self._og_logger.exchange_type = coin_enum.Okex

    self._rest_rate_limiter = RestRateLimiter(delay=10)
    self._og_logger.exchange_type = coin_enum.Okex
    self._og_logger.market_type = coin_enum.Spot

  def start(self, queue, ioloop, flow_sub=None):
    api_host = self._og_config.rest_root_url \
        if self._og_config.HasField('rest_root_url') else None
    self._private_client = OkexNativePrivateClient(key_file=self._config.key_file,
                                                   api_host=api_host,
                                                   use_async=True,
                                                   rest_lib=self._og_config.rest_lib)

    topic_string = 'okex_spot_rest_private_1'
    topic_id = topic_map.search_topic_id(topic_string)[0][1]
    self._ioloop = ioloop
    self._publisher = queue.get_writer(topic_id=topic_id)
    self._scheduler = Scheduler(ioloop=ioloop, logger=self._logger)
    self._prepare_private_subscriber()
    self._install_periodic_query()

    # Subscribe to flow for feed.
    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe(topic_string, self.on_private_exchange_msg)

  def _prepare_private_subscriber(self):
    if self._private_subscriber is not None:
      self._private_subscriber.stop()

    def on_ready():
      products = [prod.native_symbol_v3 for prod in self._og_info.products]
      self._private_subscriber.prepare_spot_subscription(products)

    self._private_subscriber = OkexWsSpotPrivateClient(
        key_file=self._config.key_file,
        ioloop=self._ioloop,
        notification_callback=self._notification_callback,
        close_callback=functools.partial(self._prepare_private_subscriber),
        on_ready=on_ready,
    )
    self._private_subscriber.start()

  def submittable(self):
    return self.is_ready() and self._rest_rate_limiter.is_available()

  @tornado.gen.coroutine
  def _install_periodic_query(self):
    assert self._scheduler is not None, 'self._scheduler is not set!'
    assert self._publisher is not None, 'self._publisher is not set!'

    query_balance_task = wrap_async_task(
        self._private_client.query_accounts,
        msg_type=self._MsgType.REST_ACCOUNT_BALANCE,
        publisher=self._publisher,
    )

    # query balance
    self._scheduler.add_periodic_task(period=self._og_config.account_query_period,
                                      job_func=query_balance_task)

    # query active orders
    for product in self._og_info.products:
      yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
      query_orders_task = wrap_async_task(
          self._private_client.query_open_order_list,
          msg_type=self._MsgType.OPEN_ORDERS,
          publisher=self._publisher,
          func_params=dict(instrument_id=product.native_symbol_v3),
      )
      self._scheduler.add_periodic_task(
          period=self._og_config.orders_snapshot_query_period,
          job_func=query_orders_task,
      )

    # query fills
    for product in self._og_info.products:
      yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
      query_fills_task = wrap_async_task(
          self._private_client.query_order_list,
          msg_type=self._MsgType.REST_FILLS,
          publisher=self._publisher,
          func_params=dict(instrument_id=product.native_symbol_v3, status=[OkexOrderStatus.FILLED]),
      )
      self._scheduler.add_periodic_task(
          period=self._og_config.fill_query_period,
          job_func=query_fills_task,
      )

      yield tornado.gen.sleep(self._og_config.burst_query_avoid_time)
      query_fills_task = wrap_async_task(
          self._private_client.query_order_list,
          msg_type=self._MsgType.REST_FILLS,
          publisher=self._publisher,
          func_params=dict(instrument_id=product.native_symbol_v3,
                           status=[OkexOrderStatus.PART_FILLED]),
      )
      self._scheduler.add_periodic_task(
          period=self._og_config.fill_query_period,
          job_func=query_fills_task,
      )

    self.start_scheduler()

  @tornado.gen.coroutine
  def _on_hit_rate_limit(self):
    self._scheduler.stop()
    yield tornado.gen.sleep(10)
    self._scheduler.start()

  def _handle_error_response(self, response):
    if response.status_code == requests.codes.too_many_requests:  # 429
      self._logger.error('Rate limited hit!')
      self._rest_rate_limiter.set_error(response.status_code)
      self._on_hit_rate_limit()
    else:
      self._logger.error('Submit response error: %s, %s', response.status_code, response.content)

  def _async_submit_order(self, order, post_only=None):
    assert order.order_id is not None, order
    assert order.order_type == OrderType.LIMIT, order
    order.tag = 'pl' + hashlib.md5(str(uuid.uuid1()).encode()).hexdigest()[:8]
    if post_only:
      order_property = OkexOrderProperty.POST_ONLY
    else:
      order_property = None

    formatter = self._order_formatter[order.product]
    params = {
        'instrument_id': order.product.native_symbol,
        'price': formatter.format_price(order.price),
        'size': formatter.format_qty(order.qty),
        'order_side': OkexOrderSide(order.order_side.name.lower()),
        'order_type': OkexOrderType(order.order_type.name.lower()),
        'order_property': order_property,
        'client_oid': order.tag,
    }
    fut = self._private_client.place_order(**params)
    return fut

  def _async_cancel_order(self, order):
    assert order.internal.exchange_order_id is not None, order
    fut = self._private_client.cancel_order(
        order_id=order.internal.exchange_order_id,
        instrument_id=order.product.native_symbol,
    )
    return fut

  def _handle_submit_response(self, order, response):
    """
    {
    "client_oid": "hello",
    "order_id": "1680287676308480",
    "result": true
    }
    """
    msg = json.loads(response.content)
    if response.status_code != requests.codes.ok:
      self._logger.error('Order submission error: %s\n%s', msg, order)
      if order in self._order_manager:
        self._order_manager.remove_order(order, ignore_error=True)
      self._handle_error_response(response)
      return OrderEvent.ORDER_REJECTED

    if order.tag != msg['client_oid']:
      self._logger.error('Weird! Inconsistent client oid!%s,%s', order.order_id, msg)
      self._order_manager.remove_order(order, ignore_error=True)
      return OrderEvent.ORDER_ERROR

    if msg['error_code'] not in (0, '', '0'):
      self._logger.error('Order submit error! %s, %s\n%s', order.order_id, msg, order)
      self._order_manager.remove_order(order, ignore_error=True)
      return OrderEvent.ORDER_ERROR

    exchange_order_id = str(msg['order_id'])
    prev_order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if prev_order is None:
      self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
    else:
      if prev_order.order_id != order.order_id:
        self._logger.info('order id replaced: %s => %s', prev_order.order_id, order.order_id)
        self._order_manager.remove_order(prev_order)
        self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
        self._order_manager.remove_order(order)
        prev_order.order_id = order.order_id
        self._order_manager.add_order(prev_order)
    return OrderEvent.ORDER_ACCEPTED

  def _handle_cancel_response(self, order, response):
    """
    {
    "client_oid": {},
    "order_id": "1680287676308480",
    "result": true
    }
    """
    msg = json.loads(response.content)
    if response.status_code != requests.codes.ok:
      if msg.get('code') == 33027:
        self._logger.error('Order cancelled already: %s, %s', msg, order)
      elif msg.get('code') == 33014:
        self._logger.error('Repeated cancellation or nonexistent order: %s, %s', msg, order)
      else:
        self._logger.error('Cancel order error! %s, %s', msg, order)
        order.cancel_sent = False
      return OrderEvent.CANCEL_ERROR

    exchange_order_id = str(msg['order_id'])
    if exchange_order_id != order.internal.exchange_order_id:
      self._logger.error('Inconsistent exchange_order_id, %s, %s', order, msg)
      order.cancel_sent = False
      return OrderEvent.CANCEL_ERROR

    if not self._order_manager.is_dead_order(order):
      self._logger.debug('%s is cancel confirmed.', order.order_id)
      self._order_manager.set_order(order, cancel_confirmed=True)
    else:
      self._logger.debug('Order is dead already %s', order)
    return OrderEvent.CANCEL_CONFIRMED

  def _notification_callback(self, data):
    timestamp = get_timestamp()
    channel = data['table']
    if channel == 'spot/account':
      msg_type = self._MsgType.WS_ACCOUNT_BALANCE
    elif channel == 'spot/order':
      msg_type = self._MsgType.ORDER_UPDATE
    else:
      self._logger.error('Unknown message type! %s', data)
      return

    exchange_message = PrivateExchangeMessage(msg_type=msg_type,
                                              data=data,
                                              received_timestamp=timestamp,
                                              request_params=None,
                                              request_timestamp=None)
    self._publisher.write(timestamp, exchange_message)

  @exchange_message_handler(_MsgType.REST_ACCOUNT_BALANCE)
  def _handle_rest_account_balance(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      self._logger.error('Rest account response error! %s, %s',
                         response.status_code,
                         response.content)
      self._handle_error_response(response)
      return

    message = json.loads(response.content)
    self._og_info.update_balance(message)
    self._og_logger.write_balance(self._og_info.balance_as_proto())
    if not self.is_ready():
      self._og_ready.set_balance_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.OPEN_ORDERS)
  def _handle_rest_open_orders(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      if response.json()['code'] == 30036:  # No working orders
        if not self.is_ready():
          self._og_ready.set_orders_ready(True)
          self._og_ready.check_ready()
      else:
        self._logger.error('Rest open orders response error! %s, %s',
                           response.status_code,
                           response.content)
      return

    timestamp = get_timestamp()
    message = response.json()

    exchange_open_order_dict = {}
    for order_msg in message:
      native_symbol = order_msg['instrument_id']
      product = OkexProduct.FromStrNativeProduct(native_symbol)
      if product not in self._og_info.products:
        continue

      order_info = create_okex_order_from_order_update(order_msg, timestamp)
      if order_info.internal.status in (OkexOrderStatus.CANCELED,
                                        OkexOrderStatus.FILLED,
                                        OkexOrderStatus.FAILURE):
        continue

      order = self._order_manager.get(exchange_order_id=order_info.internal.exchange_order_id)
      if order:
        order_info.order_id = order.order_id
      else:
        order_info.order_id = self._order_id_generator.gen_order_id(is_foreign=True)
      exchange_open_order_dict[order_info.internal.exchange_order_id] = order_info
      if self.is_ready():
        self._update_order(order_info, timestamp, is_rest=True)

    self._order_manager.copy_from_exchange_open_orders(exchange_order_dict=exchange_open_order_dict,
                                                       products=self._og_info.products)

    if not self.is_ready():
      self._og_ready.set_orders_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.REST_FILLS)
  def _handle_rest_fills(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      self._logger.error('Fill query response error! %s, %s',
                         response.status_code,
                         response.content)
      return

    timestamp = get_timestamp()
    message = response.json()
    for order_msg in message:
      order_info = create_okex_order_from_order_update(order_msg, timestamp)
      order = self._order_manager.get(exchange_order_id=order_info.internal.exchange_order_id)
      if order:
        self._update_order(order_info, timestamp, is_rest=True)

  @exchange_message_handler(_MsgType.ORDER_UPDATE)
  def _handle_ws_order_update(self, exchange_message: PrivateExchangeMessage):
    timestamp = get_timestamp()
    for order_msg in exchange_message.data['data']:
      order_info = create_okex_order_from_order_update(order_msg, timestamp)
      self._update_order(order_info, timestamp, is_rest=False)

  def _update_order(self, order_info, timestamp, is_rest):
    if order_info.product not in self._og_info.products:
      self._logger.error('Irrelevant product update received! %s', order_info)
      return

    status = order_info.internal.status
    exchange_order_id = order_info.internal.exchange_order_id
    order = self._order_manager.get(exchange_order_id=exchange_order_id)
    last_fill = None
    fill_type = None
    if order is None:
      order_id = self._order_id_generator.gen_order_id(is_foreign=True)
      order_info.order_id = order_id
      self._order_manager.add_order(order_info)
      order = order_info
      if order.internal.fill.value.qty_accum_fill > 1e-10:
        last_fill = LastFill(timestamp_last_fill=timestamp,
                             price_last_fill=order.internal.fill.value.price_avg_fill,
                             qty_last_fill=order.internal.fill.value.qty_accum_fill)
        self._logger.warning('Fill of foreign order %s', order)
    else:
      last_fill = order.internal.fill.update_by_accum_fill(
          order_info.internal.fill.value.price_avg_fill,
          order_info.internal.fill.value.qty_accum_fill)
      if last_fill is not None:
        fill_type = self._guess_fill_type(order, last_fill.price_last_fill, timestamp)

    if last_fill is not None:
      if self._order_manager.is_dead_order(order):
        self._logger.error('Weird! Dead order has fill message! %s', order)

      fill_proto = ProductFillElement(price=last_fill.price_last_fill,
                                      qty=last_fill.qty_last_fill,
                                      symbol=order.product.symbol,
                                      order_id=order.order_id,
                                      side=to_proto_order_side(order.order_side),
                                      guess_fill_type=fill_type,
                                      exchange_order_id=order.internal.exchange_order_id)
      self.process_fill(fill_proto, timestamp)
      if is_rest:
        self._logger.info('REST FILL is faster! %s', order)

    if status == OkexOrderStatus.CANCELED:
      if not self._order_manager.is_dead_order(order):
        self._order_manager.set_order(order, cancel_confirmed=True)
    elif status in (OkexOrderStatus.PART_FILLED, OkexOrderStatus.FILLED):
      if order.internal.fill.fully_filled:
        if not self._order_manager.is_dead_order(order):
          self._order_manager.set_order(order, fully_filled=True)
    elif status == OkexOrderStatus.OPEN:
      self._logger.debug('order status: %s, order_id=%s', status.name, order.order_id)
    else:
      self._logger.info('order status: %s, order_id=%s', status.name, order.order_id)
